# import torchvision
# from torch import nn
# import torch

# class ResNet(nn.Module):
#     def __init__(self, num_classes=2,model1 =torchvision.models.resnet50(pretrained=True),model2 =torchvision.models.resnet50(pretrained=True) ):
#         super(ResNet, self).__init__()

#         self.dropout = nn.Dropout()
#         self.fc = nn.Linear(1000 *2 , num_classes)
#         self.model1 = model1
#         self.model2 = model2

#     def forward(self, image , mask):
#         feature_image = self.model1(image)
#         feature_mask = self.model2(mask)
#         feature = torch.cat([feature_image, feature_mask], dim=1)
#         feature = self.dropout(feature)
#         return self.fc(feature)
from model.double.resnet import resnet18, resnet50
from model.double.senet import se_resnet18, se_resnet50
import torchvision
from torch import nn
import torch


class ResNet_double(nn.Module):
    def __init__(self, num_classes=2, model1=resnet50(pretrained=True), model2=resnet50(pretrained=True),
                 ratio=[0.8, 0.2]):
        super(ResNet_double, self).__init__()

        self.dropout = nn.Dropout()
        self.fc = nn.Linear(1000 * 2, num_classes)
        self.model1 = model1
        self.model2 = model2
        self.ratio = ratio

    def forward(self, image, mask):
        feature1_image = self.model1(image)
        feature1_mask = self.model2(mask)

        for i in range(5):
            feature1_image_cache = feature1_image
            feature1_image = torch.add(feature1_image * self.ratio[0], feature1_mask * self.ratio[1])
            feature1_mask = torch.add(feature1_image_cache * self.ratio[0], feature1_mask * self.ratio[1])
            feature1_image = self.model1(feature1_image)
            feature1_mask = self.model2(feature1_mask)

        # feature = torch.add (feature1_image,feature1_mask)
        feature = torch.cat([feature1_image, feature1_mask], dim=1)
        # feature = self.dropout(feature)
        return self.fc(feature)
